﻿#include "Data.h"
#include "utils.h"

#include <fstream>
#include <iostream>
#include <sstream>
#include <string>

#include "Layer.h"
#include "Net.h"
#include "Options.h"

#include "json.hpp"
Data::Data(Net *net, const std::string &data_path)
{
    m_net         = net;
    m_input_index = 0;

    m_default_batch           = 8;
    m_default_subdivisions    = 1;
    m_default_max_epochs      = 1000;
    m_default_hue             = 0.1f;
    m_default_saturation      = 1.5f;
    m_default_exposure        = 1.5f;
    m_default_jitter          = 0.1f; // jitter = 0.1 epochs >= 2000
    m_default_filp_horizontal = false;
    m_default_filp_vertical   = false;
    m_default_distort         = true;
    m_default_random          = true;
    m_default_auto_stop_train = true;
    m_default_debug           = false;
    // m_default_policy          = Net::s_policy_constant;
    m_default_learning_rate   = 0.0025f;
    m_default_power           = 2.0f;
    m_default_gamma           = 1.0f;
    m_default_momentum        = 0.9f;
    m_default_decay           = 0.0005f;
    std::srand(std::time(0));

    if (!data_path.empty()) Load(data_path);
}

bool Data::Load(const std::string &data_path)
{
    m_path.clear();

    auto options_list = Options::Load(data_path);
    if (options_list == nullptr || options_list->size() == 0) return false;
    auto options_list_size = options_list->size();
    for (int i = 1; i < options_list_size; i++)
    {
        auto options      = options_list->at(i);
        auto options_name = options->GetName();
        if (options_name == "default")
        {
            m_default_options = options;
        }
        else if (options_name == "train")
        {
            m_train_options = options;
        }
        else if (options_name == "valid")
        {
            m_valid_options = options;
        }
        else if (options_name == "test")
        {
            m_test_options = options;
        }
        else
        {
            std::cerr << "unknown options name: " << options_name << std::endl;
        }
    }

    auto options          = options_list->at(0);
    m_net_path            = options->GetOption("net", "");
    m_models_path         = options->GetOption("models", "");
    m_debugs_path         = options->GetOption("debugs", "");
    m_trains_path         = options->GetOption("train", "");
    m_tests_path          = options->GetOption("test", "");
    m_names_path          = options->GetOption("names", "");
    m_pretrain_model_path = options->GetOption("pretrain_model", "");

    auto data_directory = std::filesystem::path(data_path).parent_path();
    if (m_models_path.empty()) m_models_path = (data_directory / "models").string();
    if (m_debugs_path.empty()) m_debugs_path = (data_directory / "debugs").string();

    if (!LoadLines(m_names_path, &m_names))
    {
        std::cerr << "无效类别配置路径: " << m_names_path << std::endl;
        return false;
    }

    if (!LoadLines(m_trains_path, &m_trains))
    {
        std::cerr << "无效训练数据路径: " << m_trains_path << std::endl;
        return false;
    }

    if (!LoadLines(m_tests_path, &m_tests))
    {
        std::cerr << "无效训练数据路径: " << m_tests_path << std::endl;
    }

    if (!m_models_path.empty() && !std::filesystem::exists(m_models_path)) std::filesystem::create_directories(m_models_path);
    if (!m_debugs_path.empty() && !std::filesystem::exists(m_debugs_path)) std::filesystem::create_directories(m_debugs_path);

    m_path = data_path;
    return true;
}

void Data::SaveNetOptionsToDefaultValues(Options *options)
{
    LoadOptions(nullptr);
    LoadOptions(options);
    m_default_batch           = m_batch;
    m_default_subdivisions    = m_subdivisions;
    m_default_max_epochs      = m_max_epochs;
    m_default_hue             = m_hue;
    m_default_saturation      = m_saturation;
    m_default_exposure        = m_exposure;
    m_default_jitter          = m_jitter;
    m_default_filp_horizontal = m_filp_horizontal;
    m_default_filp_vertical   = m_filp_vertical;
    m_default_distort         = m_distort;
    m_default_random          = m_random;
    m_default_auto_stop_train = m_auto_stop_train;
    m_default_debug           = m_debug;
    m_default_momentum        = m_momentum;
    m_default_decay           = m_decay;
    m_default_learning_rate   = m_learning_rate;
    m_default_policy          = m_policy;
    m_default_power           = m_power;
    m_default_gamma           = m_gamma;
    m_default_steps           = m_steps;
    m_default_scales          = m_scales;
}

void Data::LoadOptions(Options *options)
{
    if (options == nullptr)
    {
        m_batch           = m_default_batch;
        m_subdivisions    = m_default_subdivisions;
        m_max_epochs      = m_default_max_epochs;
        m_hue             = m_default_hue;
        m_saturation      = m_default_saturation;
        m_exposure        = m_default_exposure;
        m_jitter          = m_default_jitter;
        m_filp_horizontal = m_default_filp_horizontal;
        m_filp_vertical   = m_default_filp_vertical;
        m_distort         = m_default_distort;
        m_random          = m_default_random;
        m_auto_stop_train = m_default_auto_stop_train;
        m_debug           = m_default_debug;
        m_momentum        = m_default_momentum;
        m_decay           = m_default_decay;
        m_policy          = m_default_policy;
        m_learning_rate   = m_default_learning_rate;
        m_power           = m_default_power;
        m_gamma           = m_default_gamma;
        m_steps           = m_default_steps;
        m_scales          = m_default_scales;
    }
    else
    {
        m_batch           = options->GetOption("batch", m_batch);
        m_subdivisions    = options->GetOption("subdivisions", m_subdivisions);
        m_max_epochs      = options->GetOption("max_epochs", m_max_epochs);
        m_jitter          = options->GetOption("jitter", m_jitter);
        m_filp_horizontal = options->GetOption("filp_horizontal", m_filp_horizontal);
        m_filp_vertical   = options->GetOption("filp_vertical", m_filp_vertical);
        m_hue             = options->GetOption("hue", m_hue);
        m_saturation      = options->GetOption("saturation", m_saturation);
        m_exposure        = options->GetOption("exposure", m_exposure);
        m_distort         = options->GetOption("distort", m_distort);
        m_random          = options->GetOption("random", m_random);
        m_auto_stop_train = options->GetOption("auto_stop_train", m_auto_stop_train);
        m_debug           = options->GetOption("debug", m_debug);
        m_momentum        = options->GetOption("momentum", m_momentum);
        m_decay           = options->GetOption("decay", m_decay);
        // m_policy          = Net::GetPolicyByName(options->GetOption("policy", ""), m_policy);
        // m_learning_rate   = options->GetOption("learning_rate", m_learning_rate);
        // m_power           = options->GetOption("power", m_power);
        // m_gamma           = options->GetOption("gamma", m_gamma);
        // if (m_policy == Net::s_policy_step)
        // {
        //     m_steps.push_back(options->GetOption("step", 1));
        //     m_scales.push_back(options->GetOption("scale", 1.0f));
        // }
        // else if (m_policy == Net::s_policy_steps)
        // {
        //     options->ParseOption("steps", "", m_steps);
        //     options->ParseOption("scales", "", m_scales);
        // }
        // else if (m_policy == Net::s_policy_sigmoid)
        // {
        //     m_steps.push_back(options->GetOption("step", 1));
        // }
    }
}

void Data::LoadDefaultOptions()
{
    LoadOptions(nullptr);
    if (m_default_options != nullptr) LoadOptions(m_default_options.get());
}

void Data::LoadTrainOptions()
{
    LoadDefaultOptions();
    if (m_train_options != nullptr) LoadOptions(m_train_options.get());
}

std::shared_ptr<Input> Data::GetNextBatchInput(std::vector<std::string> *datas)
{
    auto invalid_data_size = 0;
    auto input_item_index  = 0;
    auto net               = GetNet();
    auto datas_size        = datas->size();
    auto net_width         = net->GetWidth();
    auto net_height        = net->GetHeight();
    auto net_channels      = net->GetChannels();
    auto truth_size        = net->GetMaxTruthSize();
    auto net_batch         = GetBatch();
    auto input             = std::make_shared<Input>(net_width, net_height, net_channels, net_batch, truth_size);
    auto input_data        = input->GetData();
    auto truth_data        = input->GetTruth();
    auto net_image_size    = net_width * net_height * net_channels;
    auto is_radom          = IsRandom();
    while (net_batch > 0 && invalid_data_size < datas_size)
    {
        auto image_index = is_radom ? random_int(0, datas_size - 1) : (m_input_index % datas_size);
        auto image_path  = datas->at(image_index);

        // 加载debug信息
        auto input_item = LoadInputItem(image_path, truth_size, image_index, input_item_index);
        if (input_item == nullptr)
        {
            std::cerr << "无效数据项: " << image_path << std::endl;
            invalid_data_size++;
            m_input_index = image_index + 1;
            continue;
        }

        // 模糊图像
        input_item = AugmentInputItem(input_item);

        // 拷贝图像数据到输入对象
        memcpy(input_data, input_item->GetImage()->GetData(), net_image_size * sizeof(float));
        memcpy(truth_data, input_item->GetTruths()->data(), truth_size * sizeof(float));

        input_data += net_image_size;
        truth_data += truth_size;
        m_input_index++;
        input_item_index++;
        net_batch--;
    }

    return net_batch == 0 ? input : nullptr;
}

bool Data::LoadInputItemTruth(InputItem *input_item)
{
    std::filesystem::path path(input_item->GetImagePath());
    auto out_datas  = input_item->GetTruths()->data();
    auto layer      = GetNet()->GetBackLayer();
    auto layer_type = layer->GetType();
    if (layer_type == Layer::s_softmax)
    {
        std::stringstream ss(path.filename().string());
        int class_id = -1;
        ss >> class_id;
        if (class_id < 0 || class_id >= m_names.size() || class_id >= input_item->GetTruths()->size())
        {
            std::cerr << "分类ID不合法" << path.string() << " class_id = " << class_id << std::endl;
            return false;
        }
        out_datas[class_id] = 1.0f;
        return true;
    }
    else if (layer_type == Layer::s_yolo)
    {
        auto filename        = path.filename();
        auto label_directory = path.parent_path().parent_path() / "labels";
        auto label_path      = label_directory / filename.replace_extension(".txt");
        std::ifstream label_file(label_path);
        if (label_file.is_open())
        {
            int class_id = -1;
            label_file >> class_id;     // class_id
            label_file >> *out_datas++; // x
            label_file >> *out_datas++; // y
            label_file >> *out_datas++; // w
            label_file >> *out_datas++; // h
            if (class_id < 0 || class_id >= m_names.size())
            {
                std::cerr << "分类ID不合法" << label_path.string() << " class_id = " << class_id << std::endl;
                return false;
            }
            for (int i = 1; i <= 4; i++)
            {
                float v = *(out_datas - i);
                if (v <= 0 || v > 1)
                {
                    std::cerr << "预测框标注信息无效, 应为区间值(0, 1]" << label_path.string() << " class_id = " << class_id << std::endl;
                    return false;
                }
            }
            *out_datas++ = class_id;
            return true;
        }
        else
        {
            std::cerr << label_path.string() << "标签文件不存在" << std::endl;
        }
    }

    return false;
}

std::shared_ptr<InputItem> Data::AugmentInputItem(std::shared_ptr<InputItem> input_item)
{
    auto net           = GetNet();
    auto clone         = input_item->Clone();
    auto img           = clone->GetImage();
    auto is_augment    = net->IsTrain();
    auto is_yolo_input = GetNet()->GetBackLayer()->GetType() == Layer::s_yolo;

    // jitter
    if (m_jitter != 0 && is_augment)
    {
        auto ow = img->GetWidth();
        auto oh = img->GetHeight();

        int dw = ow * m_jitter;
        int dh = oh * m_jitter;

        int pleft  = random_int(-dw, dw);
        int pright = random_int(-dw, dw);
        int ptop   = random_int(-dh, dh);
        int pbot   = random_int(-dh, dh);

        int swidth  = ow - pleft - pright;
        int sheight = oh - ptop - pbot;

        img = img->GetSubImage(pleft, ptop, swidth, sheight);

        auto truths = clone->GetTruths()->data();
        while (is_yolo_input && truths[0] != 0.f && truths[1] != 0.f && truths[2] != 0.f && truths[3] != 0.f)
        {
            truths[0] = std::max(0.f, (truths[0] * ow - pleft) / swidth);
            truths[1] = std::max(0.f, (truths[1] * oh - ptop) / sheight);
            truths[2] = (truths[2] * ow) / swidth;
            truths[3] = (truths[3] * oh) / sheight;
            truths += 5;
            // // 变换后遮罩测试
            // {
            //     auto truths = clone->GetTruths()->data();
            //     auto img_path = std::filesystem::path(clone->GetImagePath());
            //     auto img_width = img->GetWidth();
            //     auto img_height = img->GetHeight();
            //     img->DrawRectangle((truths[0] - truths[2] / 2) * img_width, (truths[1] - truths[3] / 2) * img_height, truths[2] * img_width, truths[3] * img_height, 1, 0, 0);
            //     img->Save((img_path.parent_path().parent_path() / "tmps" / img_path.filename()).string());
            // }
        }
    }

    // resize
    auto net_width  = GetNet()->GetWidth();
    auto net_height = GetNet()->GetHeight();
    auto img_width  = img->GetWidth();
    auto img_height = img->GetHeight();
    img             = img->ScaleFill(net_width, net_height);
    if (is_yolo_input)
    {
        auto truths = clone->GetTruths()->data();
        ImageBoxToNetBox(truths, truths, img_width, img_height, net_width, net_height);
    }

    // flip
    auto is_flip_horizontal = (random_int() % 2) == 0;
    auto is_flip_vertical   = (random_int() % 2) == 0;
    if (m_filp_horizontal && is_flip_horizontal && is_augment)
    {
        img->FlipHorizontal();
        auto truths = clone->GetTruths()->data();
        while (is_yolo_input && truths[0] != 0.f && truths[1] != 0.f && truths[2] != 0.f && truths[3] != 0.f)
        {
            truths[0] = 1.0f - truths[0];
            truths += 5;
        }
    }
    if (m_filp_vertical && is_flip_vertical && is_augment)
    {
        img->FlipVertical();
        auto truths = clone->GetTruths()->data();
        while (is_yolo_input && truths[0] != 0.f && truths[1] != 0.f && truths[2] != 0.f && truths[3] != 0.f)
        {
            truths[1] = 1.0f - truths[1];
            truths += 5;
        }
    }

    // distort
    if (m_distort && is_augment) img->RandomDistort(m_hue, m_saturation, m_exposure);

    // 变换后遮罩测试
    // {
    //     auto truths = clone->GetTruths()->data();
    //     auto img_path = std::filesystem::path(clone->GetImagePath());
    //     img->DrawRectangle((truths[0] - truths[2] / 2) * net_width, (truths[1] - truths[3] / 2) * net_height, truths[2] * net_width, truths[3] * net_height, 1, 0, 0);
    //     img->Save((img_path.parent_path().parent_path() / "tmps" / img_path.filename()).string());
    // }

    clone->SetImage(img);
    assert(img->GetWidth() == net_width && img->GetHeight() == net_height);
    return clone;
}

std::shared_ptr<InputItem> Data::LoadInputItem(const std::string &in_path, const int truth_size, const int image_index, const int input_index)
{
    std::shared_ptr<InputItem> input_item = nullptr;
    std::ofstream *ofs                    = nullptr;
    auto it                               = m_image_index_items.find(image_index);
    if (it == m_image_index_items.end())
    {
        auto ofs_path = std::filesystem::path(m_debugs_path) / std::filesystem::path(in_path).filename().replace_extension(".txt");
        input_item    = std::make_shared<InputItem>(in_path, ofs_path.string());
        input_item->SetImageIndex(image_index);
        input_item->GetTruths()->resize(truth_size);

        if (IsDebug() && GetNet()->IsTrain())
        {
            input_item->SetOFStream(std::make_shared<std::ofstream>(input_item->GetOFSPath()));
            ofs = input_item->GetOFStream();
            if (ofs != nullptr)
            {
                (*ofs) << "image_path: " << in_path << std::endl;
                (*ofs) << "image_index: " << image_index << std::endl;
            }
        }

        // 加载真实值
        if (!LoadInputItemTruth(input_item.get())) return nullptr;

        // 加载图像
        auto img = std::make_shared<Image>();
        if (!img->Load(in_path)) return nullptr;
        input_item->SetImage(img);

        if (GetNet()->GetBackLayer()->GetType() == Layer::s_yolo)
        {
            auto train_ofs = GetTrainOFS();
            if (train_ofs != nullptr)
            {
                auto truths     = input_item->GetTruths()->data();
                auto net_width  = GetNet()->GetWidth();
                auto net_height = GetNet()->GetHeight();
                auto img_width  = img->GetWidth();
                auto img_height = img->GetHeight();
                std::vector<float> net_truths;
                net_truths.resize(truth_size);
                while (truths[0] != 0.f && truths[1] != 0.f && truths[2] != 0.f && truths[3] != 0.f)
                {
                    ImageBoxToNetBox(truths, net_truths.data(), img_width, img_height, net_width, net_height);
                    *train_ofs << "image path" << in_path << std::endl;
                    *train_ofs << "image truths scale: x = " << truths[0] - truths[1] / 2 << " y = " << truths[1] - truths[2] / 2 << " w = " << truths[2] << " h = " << truths[3] << std::endl;
                    *train_ofs << "image truths pixel: x = " << (truths[0] - truths[1] / 2) * img_width << " y = " << (truths[1] - truths[2] / 2) * img_height << " w = " << truths[2] * img_width << " h = " << truths[3] * img_height << std::endl;
                    *train_ofs << "net truths scale: x = " << net_truths[0] - net_truths[1] / 2 << " y = " << net_truths[1] - net_truths[2] / 2 << " w = " << net_truths[2] << " h = " << net_truths[3] << std::endl;
                    *train_ofs << "net truths pixel: x = " << (net_truths[0] - net_truths[1] / 2) * net_width << " y = " << (net_truths[1] - net_truths[2] / 2) * net_height << " w = " << net_truths[2] * net_width << " h = " << net_truths[3] * net_height << std::endl;
                    truths += 5;
                }
            }
        }
        m_image_index_items.insert_or_assign(image_index, input_item);
    }
    else
    {
        input_item = it->second;
        ofs        = input_item->GetOFStream();
    }
    m_input_index_items.insert_or_assign(input_index % GetBatch(), input_item);
    input_item->SetInputIndex(input_index);
    input_item->SetCount(input_item->GetCount() + 1);
    if (ofs != nullptr)
    {
        *ofs << std::endl;
        // auto index = GetNet()->GetEpochs() * GetEpochBatchSize() + input_index;
        // *ofs << "train index: " << index << " count: " << input_item->GetCount() << std::endl;
    }
    return input_item->Clone();
}

InputItem *Data::GetInputItem(int input_index)
{
    auto it = m_input_index_items.find(input_index % GetBatch());
    if (it == m_input_index_items.end()) return nullptr;
    return it->second.get();
}

std::ofstream *Data::GetInputOFS(int input_index)
{
    auto it = m_input_index_items.find(input_index % GetBatch());
    if (it == m_input_index_items.end()) return nullptr;
    return it->second->GetOFStream();
}

std::ofstream *Data::GetTestOFS()
{
    if (m_debugs_path.empty()) return nullptr;
    if (m_test_ofs != nullptr) return m_test_ofs.get();
    auto tests_ofs_path = std::filesystem::path(m_debugs_path) / "test.txt";
    m_test_ofs          = std::make_shared<std::ofstream>(tests_ofs_path);
    return m_test_ofs.get();
}

std::ofstream *Data::GetTrainOFS()
{
    if (m_debugs_path.empty()) return nullptr;
    if (m_train_ofs != nullptr) return m_train_ofs.get();
    auto train_ofs_path = std::filesystem::path(m_debugs_path) / "train.txt";
    m_train_ofs         = std::make_shared<std::ofstream>(train_ofs_path);
    return m_train_ofs.get();
}

bool Data::LoadLines(const std::string &in_path, std::vector<std::string> *out_lines)
{
    std::ifstream ifs(in_path);
    if (!ifs) return false;

    std::string line;
    while (std::getline(ifs, line))
    {
        // clang-format off
        line.erase(std::remove_if(line.begin(), line.end(), [](unsigned char c) -> bool { return std::isspace(c); }), line.end());
        // clang-format on
        if (line.empty() || line.front() == '#') continue;
        out_lines->push_back(line);
    }

    return true;
}

float Data::GetNetImageScale(int img_width, int img_height, int net_width, int net_height, int *offset_x, int *offset_y)
{
    float scale_x = net_width * 1.0f / img_width;
    float scale_y = net_height * 1.0f / img_height;
    float scale   = scale_x > scale_y ? scale_y : scale_x;

    if (offset_x != nullptr) *offset_x = (net_width - scale * img_width);
    if (offset_y != nullptr) *offset_y = (net_height - scale * img_height);

    return scale;
}

std::shared_ptr<Image> Data::LoadNetImage(const std::string &in_img_path, int in_net_width, int in_net_height, int *out_img_width, int *out_img_height)
{
    auto img = std::make_shared<Image>();
    if (!img->Load(in_img_path))
    {
        std::cerr << "文件路径无效: " << in_img_path << std::endl;
        return nullptr;
    }

    auto img_width  = img->GetWidth();
    auto img_height = img->GetHeight();

    if (out_img_width != nullptr) *out_img_width = img_width;
    if (out_img_height != nullptr) *out_img_height = img_height;

    if (in_net_width == img_width && in_net_height == img_height) return img;

    float scale           = GetNetImageScale(img_width, img_height, in_net_width, in_net_height, nullptr, nullptr);
    auto scale_img        = img->Scale(scale, scale);
    auto scale_img_width  = scale_img->GetWidth();
    auto scale_img_height = scale_img->GetHeight();

    if (in_net_width == scale_img_width && in_net_height == scale_img_height) return scale_img;

    return scale_img->Fill(in_net_width, in_net_height);
}

std::shared_ptr<Image> Data::LoadNetImage(const unsigned char *in_datas, const int in_size, int in_net_width, int in_net_height, int *out_img_width, int *out_img_height)
{
    auto img = std::make_shared<Image>();
    if (!img->Load(in_datas, in_size)) return nullptr;

    auto img_width  = img->GetWidth();
    auto img_height = img->GetHeight();

    if (out_img_width != nullptr) *out_img_width = img_width;
    if (out_img_height != nullptr) *out_img_height = img_height;

    if (in_net_width == img_width && in_net_height == img_height) return img;

    float scale           = GetNetImageScale(img_width, img_height, in_net_width, in_net_height, nullptr, nullptr);
    auto scale_img        = img->Scale(scale, scale);
    auto scale_img_width  = scale_img->GetWidth();
    auto scale_img_height = scale_img->GetHeight();

    if (in_net_width == scale_img_width && in_net_height == scale_img_height) return scale_img;

    return scale_img->Fill(in_net_width, in_net_height);
}

void Data::ImageBoxToNetBox(float *img_box, float *net_box, int img_width, int img_height, int net_width, int net_height)
{
    int offset_x = 0;
    int offset_y = 0;
    float scale  = GetNetImageScale(img_width, img_height, net_width, net_height, &offset_x, &offset_y);
    net_box[0]   = (img_box[0] * img_width * scale + offset_x * 0.5f) / net_width;
    net_box[1]   = (img_box[1] * img_height * scale + offset_y * 0.5f) / net_height;
    net_box[2]   = (img_box[2] * img_width * scale) / net_width;
    net_box[3]   = (img_box[3] * img_height * scale) / net_height;
    // scale = net_width / img_width | net_height / img_height
}

void Data::NetBoxToImageBox(float *net_box, float *img_box, int img_width, int img_height, int net_width, int net_height)
{
    int offset_x = 0;
    int offset_y = 0;
    float scale  = GetNetImageScale(img_width, img_height, net_width, net_height, &offset_x, &offset_y);
    img_box[0]   = (net_box[0] * net_width - offset_x * 0.5f) / scale / img_width;
    img_box[1]   = (net_box[1] * net_height - offset_y * 0.5f) / scale / img_height;
    img_box[2]   = net_box[2] * net_width / scale / img_width;
    img_box[3]   = net_box[3] * net_height / scale / img_height;
}
